from flask import render_template, redirect, flash, url_for
import os
import logging
from utils import allowed_image_file
from werkzeug.utils import secure_filename
from cos_utils import upload_to_cos
from embedding_utils import get_image_embedding
from milvus_utils import insert_image_vector, search_image_vectors

logger = logging.getLogger(__name__)

UPLOAD_FOLDER = os.getenv("UPLOAD_FOLDER", "uploads")
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)


# 处理文件上传的函数，参数请求对象request
def handle_image_upload(request):
    if request.method == "POST":
        if "file" not in request.files:
            flash("未选择文件")
            return redirect(request.url)
        try:
            file = request.files["file"]
            if file and allowed_image_file(file.filename):
                filename = secure_filename(file.filename)
                image_url = upload_to_cos(file, key=filename)
                logger.warning(f"图片上传成功:{image_url}")
                embedding = get_image_embedding(image_url)
                logger.warning(f"图片向量化成功，长度为{len(embedding)}")
                insert_image_vector(embedding, image_url)
                logger.warning(f"图片向量入库成功")
                flash("图片上传成功", "success")
                return redirect(url_for("upload_image"))
        except Exception as e:
            flash(f"图片上传失败:{e}", "error")
            logger.error(
                f"图片上传失败:{e}",
            )
    return render_template("upload_image.html")


def handle_image_search(request):
    results = []
    if request.method == "POST":
        print("request.files=", request.files)
        # 判断用户上传的表单里有没有file字段
        if "file" not in request.files:
            logger.warning("没有图片字段")
            flash("没有图片字段", "error")
            return redirect(url_for("search_image"))
        file = request.files["file"]
        #  判断字段里有没有填入真实图片
        if file.filename == "":
            logger.warning("上传的图片为空")
            flash("上传的图片为空", "error")
            return redirect(url_for("search_image"))
        if file and allowed_image_file(file.filename):
            filename = secure_filename(file.filename)
            image_url = upload_to_cos(file, key=filename)
            logger.warning(f"搜索图片上传成功:{image_url}")
            query_embedding = get_image_embedding(image_url)
            logger.warning(f"搜索图片向量化成功，长度为{len(query_embedding)}")
            # 开始到mulvus中进行搜索
            milvus_results = search_image_vectors(query_embedding, limit=5)
            logger.warning(f"检索完成，命中组数为{len(milvus_results)}")
            for hits in milvus_results:
                for hit in hits:
                    results.append(
                        {
                            "image_url": hit.entity.get("image_url"),
                            "score": hit.distance,
                        }
                    )

    return render_template("search_image.html", results=results)
